import pandas as pd
from typing import Literal, Callable

def round_dataframe(df: pd.DataFrame, decimals: int = 3) -> pd.DataFrame:
    """
    对DataFrame所有数值型列进行四舍五入处理，保留指定小数位。
    """
    numeric_cols = df.select_dtypes(include='number').columns
    df[numeric_cols] = df[numeric_cols].round(decimals)
    return df

def standardize_column_names(df: pd.DataFrame) -> pd.DataFrame:
    """
    去除所有列名的首尾空格并统一为字符串类型。
    """
    df.columns = [str(col).strip() for col in df.columns]
    return df

def drop_duplicate_rows(df: pd.DataFrame) -> pd.DataFrame:
    """
    去除重复的行，只保留唯一行（基于所有列值）
    """
    return df.drop_duplicates()

def drop_duplicate_columns(df: pd.DataFrame) -> pd.DataFrame:
    """
    去除重复的列，只保留唯一列。
    """
    return df.loc[:, ~df.T.duplicated()]

def handle_missing_values(
    df: pd.DataFrame,
    numeric_fill: str = 'mean',
    categorical_fill: str = 'mode'
) -> pd.DataFrame:
    """
    处理缺失值，对数值列和分类列分别处理
    
    参数:
        df: 输入数据框
        numeric_fill: 数值列填充方式，可选 'mean'（均值）, 'median'（中位数）, 'zero'（零）
        categorical_fill: 分类列填充方式，可选 'mode'（众数）, 'missing'（'未知'）
    """
    df = df.copy()
    
    # 处理数值列
    numeric_cols = df.select_dtypes(include=['int64', 'float64']).columns
    for col in numeric_cols:
        if numeric_fill == 'mean':
            df[col].fillna(df[col].mean(), inplace=True)
        elif numeric_fill == 'median':
            df[col].fillna(df[col].median(), inplace=True)
        elif numeric_fill == 'zero':
            df[col].fillna(0, inplace=True)
    
    # 处理分类列
    categorical_cols = df.select_dtypes(include=['object', 'category']).columns
    for col in categorical_cols:
        if categorical_fill == 'mode':
            df[col].fillna(df[col].mode()[0], inplace=True)
        elif categorical_fill == 'missing':
            df[col].fillna('未知', inplace=True)
    
    return df

def remove_outliers(
    df: pd.DataFrame,
    columns: list,
    method: str = 'iqr',
    threshold: float = 1.5
) -> pd.DataFrame:
    """
    检测和移除异常值
    
    参数:
        df: 输入数据框
        columns: 需要处理的列名列表
        method: 异常值检测方法，'iqr'（四分位距）或 'zscore'（Z分数）
        threshold: IQR法的倍数阈值或Z分数法的标准差倍数
    """
    df = df.copy()
    
    for col in columns:
        if method == 'iqr':
            Q1 = df[col].quantile(0.25)
            Q3 = df[col].quantile(0.75)
            IQR = Q3 - Q1
            lower_bound = Q1 - threshold * IQR
            upper_bound = Q3 + threshold * IQR
            df = df[(df[col] >= lower_bound) & (df[col] <= upper_bound)]
        
        elif method == 'zscore':
            z_scores = (df[col] - df[col].mean()) / df[col].std()
            df = df[abs(z_scores) <= threshold]
    
    return df

def normalize_data(
    df: pd.DataFrame,
    columns: list,
    method: str = 'minmax'
) -> pd.DataFrame:
    """
    数据归一化/标准化
    
    参数:
        df: 输入数据框
        columns: 需要处理的列名列表
        method: 归一化方法，可选 'minmax'（最小最大归一化）, 'zscore'（Z分数标准化）
    """
    df = df.copy()
    
    for col in columns:
        if method == 'minmax':
            min_val = df[col].min()
            max_val = df[col].max()
            df[col] = (df[col] - min_val) / (max_val - min_val)
        
        elif method == 'zscore':
            mean_val = df[col].mean()
            std_val = df[col].std()
            df[col] = (df[col] - mean_val) / std_val
    
    return df

def calculate_statistics(df: pd.DataFrame, columns: list) -> pd.DataFrame:
    """
    计算基本统计量
    """
    stats = []
    for col in columns:
        stats.append({
            '列名': col,
            '均值': df[col].mean(),
            '中位数': df[col].median(),
            '标准差': df[col].std(),
            '最小值': df[col].min(),
            '最大值': df[col].max(),
            '缺失值数': df[col].isnull().sum(),
            '唯一值数': df[col].nunique()
        })
    return pd.DataFrame(stats)

def bin_continuous_data(
    df: pd.DataFrame,
    column: str,
    n_bins: int = 5,
    method: str = 'equal_width'
) -> pd.Series:
    """
    连续数据分箱
    
    参数:
        df: 输入数据框
        column: 需要分箱的列名
        n_bins: 分箱数量
        method: 分箱方法，可选 'equal_width'（等宽）或 'equal_freq'（等频）
    返回:
        分箱结果，一个与输入长度相同的Series，每个值表示所属的箱子编号（从0开始）
    """
    if method == 'equal_width':
        result = pd.cut(df[column], bins=n_bins, labels=False)
    else:  # equal_freq
        result = pd.qcut(df[column], q=n_bins, labels=False)
    return pd.Series(result, index=df.index)
    
if __name__ == "__main__":

    from logger import SimpleLogger
    logger = SimpleLogger("common/test.txt")
    logger.start()
    # 构造模拟数据
    data = {
        "A": [1.12345, 2.23456, 2.23456, 4.56789, 1.12345],
        "B": [5.67891, 6.78912, 6.78912, 8.12345, 5.67891],
        "C": [9.87654, 0.12345, 0.12345, 2.34567, 9.87654],
        "D": [1.12345, 2.23456, 2.23456, 4.56789, 1.12345],
    }
    df = pd.DataFrame(data)
    df.index = pd.Index(['Row-1', 'Row-2', 'Row-3', 'Row-4', 'Row-5'])
    df.columns = [" A ", "B", "C", "D"]  # 刻意加空格

    print("原始数据：")
    print(df)

    df1 = df.copy()
    # 测试标准化列名
    df1 = standardize_column_names(df.copy())
    print("\n标准化列名：")
    print(df1.columns.tolist())

    # 测试四舍五入
    df2 = round_dataframe(df1.copy(), decimals=3)
    print("\n保留三位小数：")
    print(df2)

    print("\ndf2 的列名是:")
    print(df2.columns.to_list())

    # 测试去除重复行
    df4 = drop_duplicate_rows(df2.copy())
    print("\n去除重复行：")
    print(df4)

    # 测试去除重复列
    df5 = drop_duplicate_columns(df2.copy())
    print("\n去除重复列：")
    print(df5)

    # 测试新增的数据处理函数
    print("\n测试新增的数据处理函数：")
    
    # 创建带有缺失值的测试数据
    test_data = {
        'num1': [1, 2, None, 4, 5],
        'num2': [10, 20, 30, None, 50],
        'cat1': ['A', None, 'B', 'A', 'C'],
        'cat2': ['X', 'Y', None, 'X', 'Z']
    }
    test_df = pd.DataFrame(test_data)
    
    # 测试处理缺失值
    print("\n处理缺失值：")
    filled_df = handle_missing_values(test_df)
    print(filled_df)
    
    # 测试异常值处理
    numeric_data = {
        'values': [1, 2, 3, 100, 4, 5, -50, 6]
    }
    outlier_df = pd.DataFrame(numeric_data)
    print("\n处理异常值：")
    clean_df = remove_outliers(outlier_df, ['values'])
    print(clean_df)
    
    # 测试数据归一化
    print("\n数据归一化：")
    norm_df = normalize_data(outlier_df, ['values'])
    print(norm_df)
    
    # 测试统计量计算
    print("\n基本统计量：")
    stats_df = calculate_statistics(test_df, ['num1', 'num2'])
    print(stats_df)
    
    # 测试数据分箱
    print("\n数据分箱：")
    continuous_data = pd.DataFrame({'value': range(100)})
    binned_data = bin_continuous_data(continuous_data, 'value', n_bins=5)
    print(binned_data.value_counts().sort_index())